import anthropic
from typing import List, Dict, Any, Optional
from contentforge.core.model_interface import AIModelInterface, ModelConfig

class ClaudeModel(AIModelInterface):
    def __init__(self, config: ModelConfig):
        super().__init__(config)
        
        # 配置Anthropic客户端
        self.client = anthropic.Anthropic(
            api_key=config.api_key,
            base_url=config.api_base
        )
    
    def generate_text(
        self, 
        messages: List[Dict[str, str]], 
        system_prompt: Optional[str] = None
    ) -> str:
        """使用Claude生成文本"""
        try:
            # 准备消息
            formatted_messages = []
            
            # 添加系统提示
            if system_prompt:
                formatted_messages.append({
                    "role": "assistant",
                    "content": system_prompt
                })
            
            # 添加用户消息
            formatted_messages.extend(messages)
            
            # 调用Claude API
            response = self.client.messages.create(
                model=self.config.model_name,
                max_tokens=self.config.max_tokens,
                messages=formatted_messages
            )
            
            return response.content[0].text.strip()
        
        except Exception as e:
            raise RuntimeError(f"Claude模型文本生成失败: {str(e)}")
    
    def generate_structured_output(
        self, 
        messages: List[Dict[str, str]], 
        output_schema: Dict[str, Any],
        system_prompt: Optional[str] = None
    ) -> Dict[str, Any]:
        """生成结构化输出"""
        try:
            # 准备结构化输出提示
            structured_prompt = f"""
            请按照以下JSON模式生成输出：
            {output_schema}
            
            原始消息：{messages}
            """
            
            formatted_messages = []
            
            # 添加系统提示
            if system_prompt:
                formatted_messages.append({
                    "role": "assistant",
                    "content": system_prompt
                })
            
            # 添加用户消息
            formatted_messages.extend(messages)
            formatted_messages.append({
                "role": "user",
                "content": structured_prompt
            })
            
            # 调用Claude API
            response = self.client.messages.create(
                model=self.config.model_name,
                max_tokens=self.config.max_tokens,
                messages=formatted_messages
            )
            
            return response.content[0].text.strip()
        
        except Exception as e:
            raise RuntimeError(f"Claude模型结构化输出失败: {str(e)}") 